##########################################################################################

##https://github.com/kassambara/survminer/issues/649
##survminer和ggplot之间有冲突，最新版的ggplot会使得风险表无法对齐，需要特定版本
## mamba install conda-forge::r-ggplot2=3.4.3 只能是3.4.3
library(survival)
library(dplyr)
library(ggplot2)
library(survminer)
library(RColorBrewer)
library(data.table)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--dat_maf_file"), type = "character") ,
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--driver_list"), type = "character"),
    make_option(c("--trunk_list"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
info_file <- "MutationInfo.combine.addMolecularSubType.rmMIX.addsurv.tsv"
dat_maf_file <- "All_use.raw.maf"
driver_list <- "All_driver.list"
trunk_list <- "CompareTrunkSMG.Time.tsv"
}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

info_file <- opt$info_file
dat_maf_file <- opt$dat_maf_file
driver_list <- opt$driver_list
trunk_list <- opt$trunk_list
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################
info <- data.frame(fread(info_file))
dat_maf_cancer <- data.frame(fread( dat_maf_file))
driver_list <- data.frame(fread( driver_list ))
trunk_list <- data.frame(fread(trunk_list))
###########################################################################################

variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

dat_maf_cancer <- subset(dat_maf_cancer,Variant_Classification %in% variant_Type)


##去除MSI样本及没有预后信息的样本
info_all <- subset(info,info$OS_month!="unknown" & info$OS_status!="unknown")
info_all <- subset(info_all,info_all$MS_Type != "MSI")
info_all$OS_month <- as.numeric(info_all$OS_month)
info_all$OS_status <- as.numeric(info_all$OS_status)

## 超过60个月还存活，则认为其为存活，作为截断值
info_all$OS_status <- ifelse(info_all$OS_month > 60 , 0 , info_all$OS_status )
info_all$OS_month <- ifelse(info_all$OS_month > 60 , 60 , info_all$OS_month )

##去除stage为unknown的样本，并将年龄，性别，stage转为数值型
info_all$Stage_num <- ifelse( info_all$Stage == "I" , 1 , info_all$Stage)
info_all$Stage_num <- ifelse( info_all$Stage == "II" , 2 , info_all$Stage_num)
info_all$Stage_num <- ifelse( info_all$Stage == "III" , 3 , info_all$Stage_num)
info_all$Stage_num <- ifelse( info_all$Stage == "IV" , 4 , info_all$Stage_num)

info_all <- subset( info_all , Stage != "unknown" & Age != "[Not Available]" & Gender != "unknown" )
info_all$Stage_num <- as.numeric(info_all$Stage_num)
info_all$Age <- as.numeric(info_all$Age)
info_all$Gender <- ifelse( info_all$Gender=="Female" , 2 , 1)

info_all$From <- factor( info_all$From)

###########################################################################################

plotSurvival <- function( input = input , legend_labs=legend_labs , legend_title = legend_title , images_name = images_name ){

  ## 多因素
  test <- summary( coxph(Surv(OS_month,OS_status) ~ gene+ Age + Gender + Stage_num + From , data=input))

  res <- data.frame(
    P=round(test$coef[1,5],5),
    HR=paste(round(test$conf.int[1,1],3)," (",round(test$conf.int[1,3],3),"-",round(test$conf.int[1,4],3),")",sep=""))
  pValue=test$coef[1,5]
  if(pValue < 0.01){
    pValue <- format( pValue , scientific = TRUE , digits = 3 )
  }else{
   pValue <- round(pValue,3)
   
  }
  HR <- paste(round(test$conf.int[1,1],2)," (",round(test$conf.int[1,3],2),"-",round(test$conf.int[1,4],2),")",sep="")

  #pValue_multi=paste0("Multivariate Cox test \nHR = ", HR , "\nP = " , pValue  )
  pValue_multi=paste0("HR = ", HR , "\nP = " , pValue  )

  ## 单因素
  test <- summary( coxph(Surv(OS_month,OS_status) ~ gene , data=input))
  res <- data.frame(
    P=round(test$coef[1,5],5),
    HR=paste(round(test$conf.int[1,1],2)," (",round(test$conf.int[1,3],2),"-",round(test$conf.int[1,4],2),")",sep=""))

  pValue=test$coef[1,5]
  if(pValue < 0.01){
    pValue <- format( pValue , scientific = TRUE , digits = 3 )
  }else{
   pValue <- round(pValue,3)
   
  }

  HR <- paste(round(test$conf.int[1,1],2)," (",round(test$conf.int[1,3],2),"-",round(test$conf.int[1,4],2),")",sep="")
  pValue_single=paste0("SingleFactor Cox : HR = ", HR , ", P = " , pValue  )

  pValue <- paste0( pValue_single , "\n" , pValue_multi )

  #print("Cox adjust:")
  #print(coxph(Surv(OS_month,OS_status) ~ gene+ Age + Gender + Stage_num + From , data=input))
  #print("Cox single:")
  #print(coxph(Surv(OS_month,OS_status) ~ gene , data=input))

  bioCol <- c(  rgb(red=2,green=100,blue=190,alpha=255,max=255) , rgb(red=179,green=34,blue=35,alpha=255,max=255) )  

  fit <- survfit(Surv(OS_month,OS_status) ~ gene , data=input)

  p <- ggsurvplot(fit, 
    data=input,
    #conf.int=T,
    pval=pValue,
    pval.size=5,
    size = 2 , 
    title=legend_title,
    legend.title="",
    legend.labs=legend_labs,
    surv.median.line = "hv", # 同时显示垂直和水平参考线
    xlab="Time(month)",
    break.time.by = 5,
    palette = bioCol,
    risk.table=T,
    cumevents=F,
    ncensor.plot = F ,
    risk.table.height=.25, 
    risk.table.y.text = FALSE
    #tables.theme = theme(axis.title.y = element_text(size = 10))
    )

  pdf(images_name , width = 6.26, height = 9.19)
  print(p,newpage = FALSE)
  dev.off()

}

for(class in c("IGC","DGC","all")){
  print(class)
  dir.create(paste0(out_path,"/",class) , recursive = T)
  if(class=="all"){
    info_use <- info_all
    }else{
    info_use <- subset(info_all,info_all$Class==class) 
    }
  ##20个SMG的情况
  for(gene in driver_list$Gene_Symbol){
    print(gene)
    info_use_SMG <- info_use
    maf_use <- subset(dat_maf_cancer,dat_maf_cancer$Hugo_Symbol==gene)
    info_use_SMG$gene <- ifelse(info_use_SMG$Tumor %in% unique(maf_use$Tumor_Sample_Barcode),1,0)
    num_mut <- paste(gene,"mutated:",nrow(subset(info_use_SMG,gene==1)),sep=" ")
    num_notmut <- paste(gene,"wild-type:",nrow(subset(info_use_SMG,gene==0)),sep=" ")
    legend_title <- paste0(gene," mutation and OS")
    legend_labs <- c(num_notmut , num_mut)
    images_name <- paste0(out_path,"/",class,'/survival_',gene,"_",class,".pdf")
    plotSurvival( input = info_use_SMG , legend_labs=legend_labs , legend_title = legend_title , images_name = images_name )
  }

##含有4个maintained/4个IM favor SMG突变和不含有的预后相比
  maintained_gene <- c("TP53","APC","PIK3CA","CDH1")
  IM_favor_gene <- c("MUC6","CFTR","BMP6","MTRR")
  gene_list <- list(maintained_gene=maintained_gene,IM_favor_gene=IM_favor_gene)
  for(gene in names(gene_list)){
    print(gene)
    if(class=="IGC" & gene=="maintained_gene" ){
      ##IGC中无CDH1突变
      gene_all <- c("TP53","APC","PIK3CA")
    }else{
      gene_all <- gene_list[[gene]]
    }
    
    if(class=="all" & gene=="maintained_gene"){
    	info_use
    	maf_use_IGC <- subset(dat_maf_cancer,dat_maf_cancer$Tumor_Sample_Barcode %in% info_use$Tumor[info_use$Class=="IGC"])
    	gene_all <- c("TP53","APC","PIK3CA")
    	maf_use_IGC <- subset(maf_use_IGC,maf_use_IGC$Hugo_Symbol %in% gene_all)
    	maf_use_DGC <- subset(dat_maf_cancer,dat_maf_cancer$Tumor_Sample_Barcode %in% info_use$Tumor[info_use$Class=="DGC"])
    	gene_all <- c("TP53","APC","PIK3CA","CDH1")
    	maf_use_DGC <- subset(maf_use_DGC,maf_use_DGC$Hugo_Symbol %in% gene_all)
    	maf_use <- rbind(maf_use_IGC,maf_use_DGC)
    }else{
    	maf_use <- subset(dat_maf_cancer,dat_maf_cancer$Hugo_Symbol %in% gene_all)
    }
    info_use_maintained <- info_use
    info_use_maintained$gene <- ifelse(info_use_maintained$Tumor %in% unique(maf_use$Tumor_Sample_Barcode),1,0)
    num_mut <- paste(gene,"mutated:",nrow(subset(info_use_maintained,gene==1)),sep=" ")
    num_notmut <- paste(gene,"wild-type:",nrow(subset(info_use_maintained,gene==0)),sep=" ")
    legend_title <- paste0(gene," mutation and OS")
    legend_labs <- c(num_notmut , num_mut)
    images_name <- paste0(out_path,"/",class,'/survival_',gene,"_",class,".pdf")
    plotSurvival( input = info_use_maintained , legend_labs=legend_labs , legend_title = legend_title , images_name = images_name )
  }

  ##含有mantianed trunk突变和不含有的预后相比
  info_use_trunk <- subset(info_use,info_use$From=="NJMU")
  info_use_trunk$gene <- ifelse(info_use_trunk$Tumor %in% unique(trunk_list$ID[trunk_list$plotType=="TrunkDriver"]),1,0)
  num_mut <- paste("Trunk mutated:",nrow(subset(info_use_trunk,gene==1)),sep=" ")
  num_notmut <- paste("Trunk wild-type:",nrow(subset(info_use_trunk,gene==0)),sep=" ")
  legend_title <-paste0("Trunk mutation and OS")
  legend_labs <- c(num_notmut , num_mut)
  images_name <- paste0(out_path,"/",class,"/survival_trunk_",class,".pdf")
  plotSurvival( input = info_use_trunk , legend_labs=legend_labs , legend_title = legend_title , images_name = images_name )
}

